# coding:utf-8
# Author : hiicy redldw
# Date : 2019/03/23
from tkinter import *

# 存储树的数据结构 python可以用字典存储
class treeNode:
    def __init__(self,feat,val,right,left):
        featureToSplitOn = feat
        valueOfSplit = val
        rightBranch = right
        leftBranch = left
from numpy import *

def loadDataSet(fileName):      #general function to parse tab -delimited floats
    dataMat = []                #assume last column is target value
    fr = open(fileName)
    for line in fr.readlines():
        curLine = line.strip().split('\t')
        fltLine = list(map(float,curLine)) #map all elements to float()
        dataMat.append(fltLine)
    return dataMat

# 回归树部分
def binSplitDataSet(dataSet,feature,value):
    # REW:nonzero:1维时,也是元组,(索引,) >1维时,返回元组,长度代表轴数，每个值代表每个轴的索引;；解析：list(zip(*np.nonzero(n)))
    # 挑出数据集中该特征按照某值的划分 样本的划分;数组过滤
    mat0=dataSet[nonzero(dataSet[:,feature]>value)[0],:]
    mat1=dataSet[nonzero(dataSet[:,feature]<=value)[0],:]
    return mat0,mat1

# cart用于回归
# 最佳方式切分数据集和生成相应的叶节点
def regLeaf(dataSet):
    return mean(dataSet[:,-1]) # 创建叶节点

def regErr(dataSet): # 计算数据混乱度 用总方差来度量
    return var(dataSet[:,-1]) * shape(dataSet)[0] #因为这里需要返回的是总方差，所以要用均方差乘以数据集中样本的个数

def chooseBestSplit(dataSet, leafType, errType, ops):
    """
    ops:用户指定的参数，用于控制函数的停止时机。其中变量
    tolS是容许的误差下降值，
    tolN是切分的最少样本数
    """
    tolS = ops[0];tolN = ops[1]
    if len(set(dataSet[:,-1].T.tolist()[0])) == 1:
        return None,leafType(dataSet)
    m,n = shape(dataSet)
    S = errType(dataSet)
    bestS = inf;bestIndex=0;bestValue=0
    for featIndex in range(n-1):
        print(dataSet.A[:,featIndex])
        for splitVal in set(dataSet.A[:,featIndex]):
            mat0,mat1 = binSplitDataSet(dataSet,featIndex,splitVal)
            if (shape(mat0)[0] < tolN) or (shape(mat1)[0] < tolN):
                continue
            newS = errType(mat0) + errType(mat1)
            if newS < bestS:
                bestIndex = featIndex
                bestValue = splitVal
                bestS = newS
    if (S - bestS) < tolS:  # 如果误差减少不大则退出 也是预剪枝的一种条件
        return None, leafType(dataSet)
    mat0,mat1=binSplitDataSet(dataSet,bestIndex,bestValue)
    if (shape(mat0)[0] < tolN) or (shape(mat1)[0] < tolN):#如果切分出的数据集很小则退出
        return None,leafType(dataSet)
    return bestIndex,bestValue


def createTree(dataSet, leafType=regLeaf, errType=regErr, ops=(1,4)):#assume dataSet is NumPy Mat so we can array filtering
    # leafType给出建立叶节点的函数；
    # errType代表误差计算函数；
    # ops是一个包含树构建所需其他参数的元组
    feat, val = chooseBestSplit(dataSet, leafType, errType, ops)#choose the best split
    if feat == None: return val #if the splitting hit a stop condition return val
    retTree = {} # REW: 利用数据结构的典型算法
    retTree['spInd'] = feat
    retTree['spVal'] = val
    lSet, rSet = binSplitDataSet(dataSet, feat, val)
    retTree['left'] = createTree(lSet, leafType, errType, ops)
    retTree['right'] = createTree(rSet, leafType, errType, ops)
    return retTree

def isTree(obj):
    return (type(obj).__name__ == "dict")

def getMean(tree):
    if isTree(tree['right']):tree['right'] = getMean(tree['right'])
    if isTree(tree['left']):tree['left'] = getMean(tree['left'])
    return (tree['left']+tree['right'])/2.0

# REW:后剪枝  tree pruning 技术
def prune(tree,testData):
    if shape(testData)[0] == 0:return getMean(tree) # 没有测试数据则对树进行塌陷处理
    if (isTree(tree['right']) or isTree(tree['left'])):
        lSet,rSet = binSplitDataSet(testData,tree['spInd'],tree['spVal'])
    if isTree(tree['left']):tree['left'] = prune(tree['left'],lSet)
    if isTree(tree['right']):tree['right'] = prune(tree['right'],rSet)
    if not isTree(tree['left']) and not isTree(tree['right']):
        lSet, rSet = binSplitDataSet(testData, tree['spInd'], tree['spVal'])

        errorNoMerge = sum(power(lSet[:,-1] - tree['left'],2)) + \
                    sum(power(rSet[:,-1] - tree['right'],2))
        treeMean = (tree['left'] + tree['right']) / 2.0
        errorMerge = sum(power(testData[:,-1] - treeMean,2))
        if errorMerge < errorNoMerge:
            print('merging')
            return treeMean
        else:return tree
    else:return tree


# myDat1=loadDataSet(r'F:\Resources\Dataset\ex0_1.txt')
# myDat1=mat(myDat1)
# print(myDat1,type(myDat1))
# print(createTree(myDat1))

#模型树树仍采用二元切分，但叶节点不再是简单的数值，取而代之的是一些线性模型
def linearSolve(dataSet):   #helper function used in two places
    m,n = shape(dataSet)
    X = mat(ones((m,n))); Y = mat(ones((m,1)))#create a copy of data with 1 in 0th postion
    X[:,1:n] = dataSet[:,0:n-1]; Y = dataSet[:,-1]#and strip out Y
    xTx = X.T*X
    if linalg.det(xTx) == 0.0:
        raise NameError('This matrix is singular, cannot do inverse,\n\
        try increasing the second value of ops')
    ws = xTx.I * (X.T * Y)
    return ws,X,Y
def modelLeaf(dataSet):#create linear model and return coeficients
    ws,X,Y = linearSolve(dataSet)
    return ws
def modelErr(dataSet):
    ws,X,Y = linearSolve(dataSet)
    yHat = X * ws
    return sum(power(Y - yHat,2))
# 比较
def regTreeEval(model,inDat):
    return float(model)
def modelTreeEval(model,inDat):
    n=shape(inDat)[1]
    X = mat(ones((1,n+1)))
    X[:,1:n+1] = inDat
    return float(X*model)
def treeForeCast(tree,inData,modelEval=regTreeEval):
    if not isTree(tree):return modelEval(tree,inData)
    if inData[tree['spInd']]>tree['spVal']:
        if isTree(tree['left']):
            return treeForeCast(tree['left'],inData,modelEval)
        else:
            return modelEval(tree['left'],inData)
    else:
        if isTree(tree['right']):
            return treeForeCast(tree['right'],inData,modelEval)
        else:
            return modelEval(tree['right'],inData)
def createForeCast(tree,testData,modelEval=regTreeEval):
    m=len(testData)
    yHat = mat(zeros((m,1)))
    for i in range(m):
        yHat[i,0] = treeForeCast(tree,mat(testData[i]),modelEval)
    return yHat

def MainGui():
    import matplotlib
    matplotlib.use('TkAgg')
    from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
    from matplotlib.figure import Figure
    def reDraw(tolS, tolN):
        reDraw.f.clf()  # clear the figure
        reDraw.a = reDraw.f.add_subplot(111)
        if chkBtnVar.get():
            if tolN < 2: tolN = 2
            myTree =createTree(reDraw.rawDat, modelLeaf, \
                                         modelErr, (tolS, tolN))
            yHat = createForeCast(myTree, reDraw.testDat, \
                                           modelTreeEval)
        else:
            myTree = createTree(reDraw.rawDat, ops=(tolS, tolN))
            yHat = createForeCast(myTree, reDraw.testDat)
        reDraw.a.scatter(reDraw.rawDat[:, 0], reDraw.rawDat[:, 1], s=5)  # use scatter for data set
        reDraw.a.plot(reDraw.testDat, yHat, linewidth=2.0)  # use plot for yHat
        reDraw.canvas.show()

    def getInputs():
        try:
            tolN = int(tolNentry.get())
        except:
            tolN = 10
            print("enter Integer for tolN")
            tolNentry.delete(0, END)
            tolNentry.insert(0, '10')
        try:
            tolS = float(tolSentry.get())
        except:
            tolS = 1.0
            print("enter Float for tolS")
            tolSentry.delete(0, END)
            tolSentry.insert(0, '1.0')
        return tolN, tolS

    def drawNewTree():
        tolN, tolS = getInputs()  # get values from Entry boxes
        reDraw(tolS, tolN)

    root = Tk()

    reDraw.f = Figure(figsize=(5, 4), dpi=100)  # create canvas
    reDraw.canvas = FigureCanvasTkAgg(reDraw.f, master=root)
    reDraw.canvas.show()
    reDraw.canvas.get_tk_widget().grid(row=0, columnspan=3)

    Label(root, text="tolN").grid(row=1, column=0)
    tolNentry = Entry(root)
    tolNentry.grid(row=1, column=1)
    tolNentry.insert(0, '10')
    Label(root, text="tolS").grid(row=2, column=0)
    tolSentry = Entry(root)
    tolSentry.grid(row=2, column=1)
    tolSentry.insert(0, '1.0')
    Button(root, text="ReDraw", command=drawNewTree).grid(row=1, column=2, rowspan=3)
    chkBtnVar = IntVar()
    chkBtn = Checkbutton(root, text="Model Tree", variable=chkBtnVar)
    chkBtn.grid(row=3, column=0, columnspan=2)

    reDraw.rawDat = mat(loadDataSet('sine.txt'))
    reDraw.testDat = arange(min(reDraw.rawDat[:, 0]), max(reDraw.rawDat[:, 0]), 0.01)
    reDraw(1.0, 10)

    root.mainloop()
MainGui()